import pwlf
import numpy as np
# import matplotlib.pyplot as plt

# 以上的变量都拿到以后 然后把上面的变量放入到下面的函数里面 输出 status
# createNozzleMap输入：
#  pixelSeq，也就是输入的喷阀对应的图片坐标的list如[5,7650]，  下面就是x
#  nuzzleSeq，输入对应的喷阀序号组成的list，如[1,256]          下面就是y
#  int(nuzzleNumTotal)   输入的喷阀总数（转为整形）            下面是同名
#  savePath 就是保存文件的路径                                   下面是同名
#  nozzleMap 就是上面的相机分辨率  天元的选择的就是7680 或者手动输入的其他方案   下面是同名，  下面的函数里面改成了pixelNum
#  curzeType 就是上面的拟合方式 暂时目前用的都是现行拟合                 下面是同名
#  int(badNozzle) 就是不可用的阀的序号（转为整形）                      下面是同名
# 以上就是这几个函数的输入的参数 后面就是去拟合了
# status = createNozzleMap(pixelSeq, nuzzleSeq, int(nuzzleNumTotal), savePath, nozzleMap, curzeType,
#                          int(badNozzle))
def createNozzleMap(x,y,nozzleNumTotal,savePath,pixelNum=8192,curzeType = True,badNozzle=0):
    #debug
    #y=[1,11,21,31,41,51,61,71,81,91,101,111,121,131,141,150]
    #x=[54,629,1197,1680,2208,2762,3334,3862,4390,4938,5517,6046,6606,7144,7680,8180]
    # y=[1,11,21,31,41,51,61,71,81,91,100]
    # x=[5456,5040,4620,4208,3792,3372,2956,2546,2124,1724,1344]
    y2=[]
    print("pixelNum:",pixelNum)#pixelNum 如果不选择 那就是默认的8192  这里选择了7680那就是7680
    xHatTotal = np.linspace(1,pixelNum,num=pixelNum)#这里最大值和要求生成的点的个数是一样的  也就是1，2,3...7680
    #curzeType == True 如果拟合方式是线性拟合
    if curzeType == True:
        # test
        # testX = [x[0],x[-1]]
        # testY = [y[0],y[-1]]
        # my_pwlf = pwlf.PiecewiseLinFit(testX, testY)
        # my_pwlf.fit_with_breaks(testX)
        
        my_pwlf = pwlf.PiecewiseLinFit(x, y)#x就是[5,7650]   y就是[1,256]  然后就放入pwlf里面进行训练 拟合
        my_pwlf.fit_with_breaks(x)#这句代码虽然不知道作用是什么  然后就按照步骤来吧
        xHat = np.linspace(1, pixelNum, num=pixelNum)#这里就生成了1,2,3,4...7680个数字
        with open( savePath+'//pixNumberNozzleMap.txt', 'wt') as f:
            if(badNozzle !=0):#如果不可用的阀的序号不是0的话 count=1
                count = 1
            for idx,s in enumerate(xHat):#开始便利xHat里面的数据
                yHat = my_pwlf.predict(s).round().astype(int)#用之前拟合好的模型来预测y
                #print(yHat)
                nozzleNum = int(yHat.item())#然后把预测的y转成整形
                if nozzleNum <= 0:#如果预测出来的y是小于0的 就直接变成0
                    nozzleNum = 0
                elif nozzleNum >= nozzleNumTotal:#如果这个预测的y大于喷阀的总数
                    if(badNozzle != 0):    #如果不可用的阀的序号不是0 也就是其他的1 2 3 之类的
                        if(nozzleNum >= badNozzle):#如果是预测的y是大于不可用的阀的序号 就让y-1
                            nozzleNum = nozzleNum-1
                        nozzleNum = nozzleNumTotal-2#预测的y的值等于喷阀总数减2
                    else:#如果不可用的阀的序号是0 这个预测的y就等于喷阀的总数就减1
                        nozzleNum = nozzleNumTotal-1

                else:#综合上面的 if和elif  就是如果预测的y 是大于0 且不大于喷阀总数  就是介于0-喷阀总数之间的话
                    if(badNozzle != 0):#如果不可用的阀的序号不是0
                        if(nozzleNum >= badNozzle):#如果预测的y大于等于不可用的阀的序号
                            if(count <= round(pixelNum/nozzleNumTotal/2)):
                                nozzleNum = nozzleNum-2
                            elif(count > round(pixelNum/nozzleNumTotal/2) and count <= round(pixelNum/nozzleNumTotal)):
                                nozzleNum = nozzleNum
                            else:
                                nozzleNum = nozzleNum-1
                            count+=1
                        else:#如果预测的y小于不可用的阀的序号
                            nozzleNum = nozzleNum-1
                    else:#如果不可用的阀的序号是0
                        nozzleNum =nozzleNum-1

                y2.append(nozzleNum)
                if idx == xHat.size-1:
                    f.write(str(nozzleNum)+'\n')
                else:
                    f.write(str(nozzleNum)+' ')
    else:#这个是分段拟合 这里暂时用不上！
        flag = 0
        if x[0] > x[-1]:
            flag = 1
            x.reverse()
            y.reverse()
        # if x[0] !=1:
        #     x.insert(0,1)
        # if x[-1] != pixelNum:
        #     x.append(pixelNum)
        
        if flag:
            if y[0] != nozzleNumTotal:
                y.insert(0,nozzleNumTotal)
                if x[0] !=1:
                    x.insert(0,1)
            if y[-1] != 1:
                y.append(1)
                if x[-1] != pixelNum:
                    x.append(pixelNum)
        else:
            if y[0] != 1:
                y.insert(0,1)
                if x[0] !=1:
                    x.insert(0,1)
            if y[-1] != nozzleNumTotal:
                y.append(nozzleNumTotal)
                if x[-1] != pixelNum:
                    x.append(pixelNum)
                    
        blockSize = len(x)-1
        with open( savePath+'//pixNumberNozzleMap.txt', 'wt') as f:

            for i in range(blockSize):
                tempX = []
                tempY = []
                tempX = [x[i],x[i+1]]
                tempY = [y[i],y[i+1]]
                my_pwlf = pwlf.PiecewiseLinFit(tempX, tempY)
                my_pwlf.fit_with_breaks(tempX)
                # if x[0]>x[-1]:
                #     xHat = np.linspace(x[i], x[i+1], num=abs(x[i+1]-x[i]) )
                # else:
                xHat = np.linspace(x[i], x[i+1], num=abs(x[i+1]-x[i]) )

                for idx,s in enumerate(xHat):
                    yHat = my_pwlf.predict(s).round().astype(int)
                    nozzleNum = int(yHat.item())
                    if nozzleNum <=0:
                        nozzleNum = 0
                    elif nozzleNum >= nozzleNumTotal:
                        nozzleNum = nozzleNumTotal -1
                    else:
                        nozzleNum -=1
                    y2.append(nozzleNum)
                    if idx == xHatTotal.size -1:
                        f.write(str(nozzleNum)+'\n')
                    else:
                        f.write(str(nozzleNum)+ ' ')

            left = xHatTotal.size - len(y2)
            
            for j in range(left):
                if(flag):
                    y2.append(0)
                    f.write(str(0)+ ' ')

                else:
                    y2.append(nozzleNumTotal-1)
                    f.write(str(nozzleNumTotal-1)+ ' ')

                # f.write(str(0)+ ' ')
    
    # plt.figure()
    #plt.plot(x, y2, 'o')
    # plt.plot(xHatTotal, y2, '-')
    # plt.show(block=False)
    
    return 1
    